import numpy as np
import torch
from scipy import io
from torch.utils.data import Dataset

from common_utils.utils import l2_normalize

device = 'cuda' if torch.cuda.is_available() else 'cpu'


class Dataset(Dataset):
    '''
    Inputs
        data_dir String 数据集根目录
        dataset String 数据集
        phase String train, val, test
    '''

    def __init__(self, data_dir, dataset, phase):
        self.data_dir = data_dir
        self.dataset = dataset
        self.phase = phase

        res101 = io.loadmat(data_dir +dataset+'/'+ 'res101.mat')
        att_splits = io.loadmat(data_dir  +dataset+'/'+ 'att_splits.mat')
        train_loc = 'train_loc'
        val_loc = 'val_loc'
        trainval_loc = 'trainval_loc'
        test_unseen_loc = 'test_unseen_loc'
        # 取出训练好的特征
        feat = res101['features'].T
        self.feat_dim = feat.shape[1]
        # 提取训练图片的位置
        self.train_data = feat[np.squeeze(att_splits[train_loc] - 1)]
        self.train_data = torch.tensor(self.train_data,dtype=torch.float32)
        self.val_data = feat[np.squeeze(att_splits[val_loc] - 1)]
        self.val_data = torch.tensor(self.val_data, dtype=torch.float32)
        self.trainval_data = feat[np.squeeze(att_splits[trainval_loc] - 1)]
        self.trainval_data = torch.tensor(self.trainval_data, dtype=torch.float32)
        self.test_unseen_data = feat[np.squeeze(att_splits[test_unseen_loc] - 1)]
        self.test_unseen_data = torch.tensor(self.test_unseen_data, dtype=torch.float32)
        print('Tr:{}; Val:{}; Ts:{}\n'.format(self.train_data.shape[0], self.val_data.shape[0], self.test_unseen_data.shape[0]))

        labels = res101['labels']
        # 获取图片位置对于的分类
        self.labels_train = np.squeeze(labels[np.squeeze(att_splits[train_loc] - 1)])
        self.labels_val = np.squeeze(labels[np.squeeze(att_splits[val_loc] - 1)])
        self.labels_trainval = np.squeeze(labels[np.squeeze(att_splits[trainval_loc] - 1)])
        self.labels_test = np.squeeze(labels[np.squeeze(att_splits[test_unseen_loc] - 1)])

        # unique函数去除其中重复的元素，并按元素由大到小返回一个新的无元素重复的元组或者列表
        self.train_classes_seen = np.unique(self.labels_train)
        self.val_classes_seen = np.unique(self.labels_val)
        self.trainval_classes_seen = np.unique(self.labels_trainval)
        self.test_classes_seen = np.unique(self.labels_test)

        train_gzsl_indices = []
        val_gzsl_indices = []

        for cl in self.train_classes_seen:
            train_gzsl_indices = np.squeeze(np.where(self.labels_trainval == cl)).tolist()

        for cl in self.val_classes_seen:
            val_gzsl_indices = np.squeeze(np.where(self.labels_trainval == cl)).tolist()

        i = 0
        for labels in self.train_classes_seen:
            self.labels_train[self.labels_train == labels] = i
            i += 1

        j = 0
        for labels in self.val_classes_seen:
            self.labels_val[self.labels_val == labels] = j
            j += 1

        k = 0
        for labels in self.test_classes_seen:
            self.labels_test[self.labels_test == labels] = k
            k += 1
        # 获取已经归一化的属性矩阵，50*85, 50个分类，每个分类85个属性
        self.sig = att_splits['att'].T
        self.attr_dim = self.sig.shape[1]
        # Shape -> (Number of attributes, Number of Classes)
        self.train_attrs = self.sig[self.train_classes_seen - 1]
        self.train_attrs = torch.tensor(self.train_attrs, dtype=torch.float32)
        self.val_attrs = self.sig[self.val_classes_seen - 1]
        self.val_attrs = torch.tensor(self.val_attrs, dtype=torch.float32)
        self.test_attrs = self.sig[self.test_classes_seen - 1]
        self.test_attrs = torch.tensor(self.test_attrs, dtype=torch.float32)
        if self.phase == 'train':
            self.data = self.train_data
            self.labels = self.labels_train
        elif self.phase == 'val':
            self.data = self.val_data
            self.labels = self.labels_val
        elif  self.phase == 'test':
            self.data = self.test_unseen_data
            self.labels = self.labels_test
        elif  self.phase == 'all':
            self.data = self.train_data + self.val_data + self.test_unseen_data
            self.labels = labels

        self.labels = torch.tensor(self.labels, dtype=torch.long)
    def __getitem__(self, index):
        img = self.data[index]
        label = self.labels[index]
        # label = torch.tensor(label, dtype=torch.long)
        attr_id = self.train_attrs[label]
        img = l2_normalize(img, dim=0)
        return img, label, attr_id

    def __len__(self):
        return len(self.data)